import numpy as np
from nets.bases.basicmodel import BasicModule
import torch as t
import torch.nn as nn
import torch.nn.functional as F
from nets.modules.depthwise_block import DepthWiseConv
from nets.modules.dilated_conv import DilatedConv
import torchvision
from matplotlib import pyplot as plt

class DownSample(nn.Module):
    
    def __init__(self, in_plances, out_plances, isDown = True):
        super().__init__()
        
        self.model = nn.Sequential(
            nn.Conv2d(in_plances, out_plances, (3, 3), stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_plances),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(out_plances, out_plances, (3, 3), stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_plances),
            nn.ReLU(inplace=True),
            
        )
        
        self.isDown = isDown

    def forward(self, imgs):
        x = self.model(imgs)
        if self.isDown:
            x = F.max_pool2d(x, 2)
        return x
    
class UpSample(nn.Module):
    
    def __init__(self, in_plances, out_plances, isUp = True):
        super().__init__()

        self.model = nn.Sequential(
            nn.Conv2d(in_plances, out_plances, (3, 3), stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_plances),
            nn.ReLU(inplace=True),
        )
        if isUp:        
            self.upSample = nn.Upsample(scale_factor=2, mode='bilinear',align_corners=False)
        else:
            self.upSample = nn.Conv2d(out_plances, out_plances, (1, 1), stride=1, padding=0, bias=False)
        
    
    def forward(self, x):
        x = self.model(x)
        x = self.upSample(x)
        return x

class UNetModel(BasicModule):
    
    def __init__(self):
        super(UNetModel, self).__init__()  # 初始化父类
        self.model_name = 'UNetModel'
        self.encoding_layer1 = DownSample(3, 64)     #512 -> 256
        self.encoding_layer2 = DownSample(64, 128)   #256 -> 128
        self.encoding_layer3 = DownSample(128, 256)  #128 -> 64
        self.encoding_layer4 = DownSample(256, 512)  #64 -> 32
        #编码最后一层
        self.encoding_layer5 = DownSample(512, 512, False) #32 -> 32
        
        self.deconding_layer1 = UpSample(1024, 256)        #32 -> 64
        self.deconding_layer2 = UpSample(512, 128)         #64 -> 128
        self.deconding_layer3 = UpSample(256, 64)          #128 -> 256
        self.deconding_layer4 = UpSample(128 + 16, 16, False)   #256 -> 256
        
        self.encoding_layer1_skip_connection = nn.Sequential(
            nn.Conv2d(3, 16, (3, 3), stride=2, padding=1, bias=False),
            # nn.MaxPool2d(2),
        )
        
        #最后一层
        self.last_layer = nn.Conv2d(16, 1, (1, 1), stride=1, padding=0, bias=False)        
        
        
        
        
    #前向传播
    def forward(self, imgs):
        x1 = self.encoding_layer1(imgs)
        x2 = self.encoding_layer2(x1)
        x3 = self.encoding_layer3(x2)
        x4 = self.encoding_layer4(x3)
        
        x5 = self.encoding_layer5(x4)
        
        x = t.cat((x5, x4), dim=1)
        x = self.deconding_layer1(x)
        x = t.cat((x, x3), dim=1)
        x = self.deconding_layer2(x)
        x = t.cat((x, x2), dim=1)
        x = self.deconding_layer3(x)
        c_x = self.encoding_layer1_skip_connection(imgs)
        x = t.cat((x, x1, c_x), dim=1)
        x = self.deconding_layer4(x)
        x = self.last_layer(x)
        # x = F.sigmoid(x)
        return x


class  DepthDownSample(nn.Module):
    
    def __init__(self, in_plances, out_plances, isDown = True):
        super().__init__()
        
        self.model = nn.Sequential(
            DepthWiseConv(in_plances, in_plances),
            nn.BatchNorm2d(in_plances),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(in_plances, out_plances, (3, 3), stride=1, padding=1, bias=False),
            # DepthWiseConv(in_plances, out_plances),
            nn.BatchNorm2d(out_plances),
            nn.ReLU(inplace=True),
        )
        
        self.isDown = isDown

    def forward(self, imgs):
        x = self.model(imgs)
        if self.isDown:
            x = F.max_pool2d(x, 2)
        return x

class DepthUNetModel(BasicModule):
    
    def __init__(self):
        super(DepthUNetModel, self).__init__()  # 初始化父类
        self.model_name = 'UNetModel'
        self.encoding_layer1 = DepthDownSample(3, 18)     #512 -> 256
        self.encoding_layer2 = DepthDownSample(18, 126)   #256 -> 128
        self.encoding_layer3 = DepthDownSample(126, 252)  #128 -> 64
        # self.encoding_layer4 = DepthDownSample(252, 504)  #64 -> 32
        
        
        # self.deconding_layer1 = UpSample(1008, 252)        #32 -> 64
        self.deconding_layer2 = UpSample(504, 126)         #64 -> 128
        self.deconding_layer3 = UpSample(252, 18)          #128 -> 256
        self.deconding_layer4 = UpSample(52, 15, False)    #256 -> 256
        
        
        self.dilated1 = DilatedConv(inplace = 252, outplace = 84, dilation_rate = 1, stride = 1, kernel_size = 1, padding = 0)
        self.dilated2 = DilatedConv(inplace = 252, outplace = 84, dilation_rate = 2, stride = 1, kernel_size = 3, padding = 2)
        self.dilated3 = DilatedConv(inplace = 252, outplace = 84, dilation_rate = 5, stride = 1, kernel_size = 3, padding = 5)
        
        self.encoding_layer1_skip_connection = nn.Sequential(
            nn.Conv2d(3, 16, (3, 3), stride=2, padding=1, bias=False),
        )
        
        #最后一层
        self.last_layer = nn.Conv2d(15, 1, (1, 1), stride=1, padding=0, bias=False)        
        
        
        
    #前向传播
    def forward(self, imgs):
        x1 = self.encoding_layer1(imgs)
        x2 = self.encoding_layer2(x1)
        x3 = self.encoding_layer3(x2)
        # x4 = self.encoding_layer4(x3)
        
        d1 = self.dilated1(x3)
        d2 = self.dilated2(x3)
        d3 = self.dilated3(x3)
        
        x5 = t.cat((d1, d2, d3), dim=1)
        
        # x = t.cat((x5, x4), dim=1)
        # x = self.deconding_layer1(x)
        x = t.cat((x3, x5), dim=1)
        x = self.deconding_layer2(x)
        x = t.cat((x, x2), dim=1)
        x = self.deconding_layer3(x)
        c_x = self.encoding_layer1_skip_connection(imgs)
        x = t.cat((x, x1, c_x), dim=1)
        x = self.deconding_layer4(x)
        x = self.last_layer(x)
        return x
    
    
    #模型可视化
    def visualize(self, imgs):
        f = plt.figure()
        self.__convert_to_map(imgs, [3, 3, 9], [1 , 3, 512, 512], title = 'layer1')
        
        x1 = self.encoding_layer1(imgs)
        x2 = self.encoding_layer2(x1)
        # 保存第二层的输出
        # print(x2.shape)
        self.__convert_to_map(x2, [3, 3, 1], [126 , 1, 128, 128], title = 'layer2')
        x3 = self.encoding_layer3(x2)
        #保存第三层的输出
        self.__convert_to_map(x3, [3, 3, 2], [252 , 1, 64, 64], visual_max_count=120, title = 'layer3')
        # x4 = self.encoding_layer4(x3)
        # #保存第四层的输出
        # self.__convert_to_map(x4, [3, 3, 3], [504, 1, 32, 32], visual_max_count=240, title = 'layer4')
        
        d1 = self.dilated1(x3)
        d2 = self.dilated2(x3)
        d3 = self.dilated3(x3)
        x5 = t.cat((d1, d2, d3), dim=1)
        
        #保存第五层的输出
        self.__convert_to_map(x5, [3, 3, 4], [252, 1, 64, 64], visual_max_count=240, title = 'layer5')
        
        # x = t.cat((x5, x4), dim=1)
        # x = self.deconding_layer1(x)
        #保存第六层的输出
        # self.__convert_to_map(x, [3, 3, 5], [252, 1, 64, 64], visual_max_count=120, title = 'layer6')
        
        x = t.cat((x3, x5), dim=1)
        x = self.deconding_layer2(x)
        #保存第七层的输出
        self.__convert_to_map(x, [3, 3, 6], [126, 1, 128, 128], title = 'layer7')
        
        x = t.cat((x, x2), dim=1)
        x = self.deconding_layer3(x)
        #保存第八层的输出
        self.__convert_to_map(x, [3, 3, 7], [18, 1, 256, 256], title = 'layer8')
        
        c_x = self.encoding_layer1_skip_connection(imgs)
        x = t.cat((x, x1, c_x), dim=1)
        x = self.deconding_layer4(x)
        #保存第九层的输出
        self.__convert_to_map(x, [3, 3, 8], [15, 1, 256, 256], title = 'layer9')
        
        x = self.last_layer(x)
        
        plt.show()
        return x

    
    #转换到可以显示的map
    def __convert_to_map(self, map : t.tensor, grid, show_shape = None, size = (100, 100), visual_max_count = 30, title = 'map'):
        map = map.detach().reshape(show_shape)
        # map = F.normalize(map, p=2, dim=1)
        img : np.ndarray = None
        if map.shape[0] > visual_max_count:
            img = map[:visual_max_count].cpu()
        else:
            img = map.cpu()
        img = torchvision.utils.make_grid(img, nrow=15).numpy()
        #不能小于0
        img = np.maximum(img, 0)
        img = img / img.max()
        img = img.transpose(1, 2, 0)
        
        
        plt.subplot(grid[0], grid[1], grid[2])
        plt.imshow(img)
        plt.axis('on')        
        plt.title(title)
        